#!/usr/bin/env python3
"""高级使用示例

展示 auto_video 项目的高级功能，包括模型管理、批量生成、图像处理等。
"""

import sys
import os
from pathlib import Path
from typing import List, Dict, Any

# 添加项目根目录到 Python 路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from src.generators.text_to_image import TextToImage
from src.models.model_manager import get_model_manager
from src.utils.config import get_config
from src.utils.logger import setup_logging, LogContext, ProgressLogger
from src.utils.image_utils import ImageProcessor, create_output_filename
from src.utils.file_utils import FileManager, ensure_output_dir


def batch_generation_example():
    """批量生成示例"""
    logger = setup_logging()
    
    with LogContext(logger, "批量图片生成"):
        # 准备批量提示词
        batch_prompts = [
            {
                "prompt": "a majestic mountain landscape at sunrise",
                "negative_prompt": "blurry, low quality",
                "width": 768,
                "height": 512,
                "steps": 25,
                "guidance": 8.0
            },
            {
                "prompt": "a cyberpunk street scene with neon lights",
                "negative_prompt": "blurry, low quality, bad anatomy",
                "width": 512,
                "height": 768,
                "steps": 30,
                "guidance": 7.5
            },
            {
                "prompt": "a serene japanese garden with cherry blossoms",
                "negative_prompt": "ugly, distorted",
                "width": 512,
                "height": 512,
                "steps": 20,
                "guidance": 7.0
            }
        ]
        
        # 创建输出目录
        output_dir = ensure_output_dir("outputs/batch_generation")
        
        # 初始化生成器
        generator = TextToImage()
        
        # 进度跟踪
        progress = ProgressLogger(logger, len(batch_prompts), "批量生成")
        
        results = []
        
        for i, config in enumerate(batch_prompts):
            try:
                # 生成图片
                image = generator.generate(
                    prompt=config["prompt"],
                    negative_prompt=config.get("negative_prompt"),
                    width=config["width"],
                    height=config["height"],
                    num_inference_steps=config["steps"],
                    guidance_scale=config["guidance"]
                )
                
                # 保存图片
                filename = create_output_filename(f"batch_{i+1:03d}", ".jpg")
                output_path = output_dir / filename
                generator.save_image(image, output_path)
                
                results.append({
                    "index": i + 1,
                    "prompt": config["prompt"],
                    "output_path": str(output_path),
                    "success": True
                })
                
                progress.update(1, f"完成: {config['prompt'][:30]}...")
                
            except Exception as e:
                logger.error(f"生成第 {i+1} 张图片失败: {e}")
                results.append({
                    "index": i + 1,
                    "prompt": config["prompt"],
                    "error": str(e),
                    "success": False
                })
                progress.update(1, f"失败: {str(e)[:30]}...")
        
        progress.finish()
        
        # 统计结果
        successful = sum(1 for r in results if r["success"])
        logger.info(f"批量生成完成: {successful}/{len(batch_prompts)} 成功")
        
        return results


def model_management_example():
    """模型管理示例"""
    logger = setup_logging()
    
    with LogContext(logger, "模型管理演示"):
        # 获取模型管理器
        model_manager = get_model_manager()
        
        # 显示可用模型
        available_models = model_manager.get_available_models()
        logger.info(f"可用模型数量: {len(available_models)}")
        
        for model in available_models:
            logger.info(f"- {model['name']}: {model['id']} ({model['size']})")
        
        # 检查缓存状态
        default_model = "runwayml/stable-diffusion-v1-5"
        is_cached = model_manager.is_model_cached(default_model)
        logger.info(f"模型 {default_model} 缓存状态: {'已缓存' if is_cached else '未缓存'}")
        
        # 获取模型信息
        model_info = model_manager.get_model_info(default_model)
        if model_info:
            logger.info(f"模型信息: 作者={model_info.get('author')}, 下载量={model_info.get('downloads')}")
        
        # 获取缓存大小信息
        cache_info = model_manager.get_cache_size()
        logger.info(f"缓存总大小: {cache_info['total_size_gb']:.2f} GB")
        
        return model_manager


def image_processing_example():
    """图像处理示例"""
    logger = setup_logging()
    
    with LogContext(logger, "图像处理演示"):
        # 创建输出目录
        output_dir = ensure_output_dir("outputs/image_processing")
        
        # 首先生成一张基础图片
        generator = TextToImage()
        base_image = generator.generate(
            prompt="a beautiful landscape with mountains and lake",
            width=512,
            height=512,
            num_inference_steps=20
        )
        
        # 保存原始图片
        original_path = output_dir / "original.jpg"
        generator.save_image(base_image, original_path)
        logger.info(f"原始图片已保存: {original_path}")
        
        # 图像处理操作
        processor = ImageProcessor()
        
        # 1. 调整尺寸
        resized = processor.resize_image(base_image, (768, 768))
        processor.save_image(resized, output_dir / "resized.jpg")
        logger.info("完成尺寸调整")
        
        # 2. 中心裁剪
        cropped = processor.center_crop(base_image, (400, 400))
        processor.save_image(cropped, output_dir / "cropped.jpg")
        logger.info("完成中心裁剪")
        
        # 3. 图像增强
        enhanced = processor.enhance_image(
            base_image,
            brightness=1.2,
            contrast=1.1,
            saturation=1.3,
            sharpness=1.1
        )
        processor.save_image(enhanced, output_dir / "enhanced.jpg")
        logger.info("完成图像增强")
        
        # 4. 应用滤镜
        blurred = processor.apply_filter(base_image, "gaussian_blur", radius=2)
        processor.save_image(blurred, output_dir / "blurred.jpg")
        logger.info("完成模糊滤镜")
        
        sharpened = processor.apply_filter(base_image, "sharpen")
        processor.save_image(sharpened, output_dir / "sharpened.jpg")
        logger.info("完成锐化滤镜")
        
        # 5. 创建图像网格
        images = [base_image, resized, cropped, enhanced, blurred, sharpened]
        grid = processor.create_grid(
            images[:4],  # 只使用前4张图片
            grid_size=(2, 2),
            image_size=(256, 256),
            spacing=10
        )
        processor.save_image(grid, output_dir / "grid.jpg")
        logger.info("完成图像网格创建")
        
        logger.info("图像处理演示完成")


def configuration_example():
    """配置管理示例"""
    logger = setup_logging()
    
    with LogContext(logger, "配置管理演示"):
        # 获取配置
        config = get_config()
        
        # 显示当前配置
        logger.info("当前配置信息:")
        logger.info(f"- 默认模型: {config.get('model.default_model_id')}")
        logger.info(f"- 默认尺寸: {config.get('generation.default_width')}x{config.get('generation.default_height')}")
        logger.info(f"- 默认步数: {config.get('generation.default_num_inference_steps')}")
        logger.info(f"- 输出目录: {config.get('output.default_output_dir')}")
        
        # 临时修改配置
        original_steps = config.get('generation.default_num_inference_steps')
        config.set('generation.default_num_inference_steps', 25)
        logger.info(f"临时修改步数: {original_steps} -> {config.get('generation.default_num_inference_steps')}")
        
        # 恢复配置
        config.set('generation.default_num_inference_steps', original_steps)
        logger.info(f"恢复步数: {config.get('generation.default_num_inference_steps')}")


def main():
    """主函数"""
    logger = setup_logging(level='INFO')
    logger.info("开始高级使用示例演示")
    
    try:
        # 1. 模型管理示例
        logger.info("=== 模型管理示例 ===")
        model_manager = model_management_example()
        
        # 2. 配置管理示例
        logger.info("\n=== 配置管理示例 ===")
        configuration_example()
        
        # 3. 批量生成示例
        logger.info("\n=== 批量生成示例 ===")
        batch_results = batch_generation_example()
        
        # 4. 图像处理示例
        logger.info("\n=== 图像处理示例 ===")
        image_processing_example()
        
        logger.info("\n=== 所有高级功能演示完成 ===")
        
        # 显示总结
        successful_batch = sum(1 for r in batch_results if r["success"])
        logger.info(f"批量生成成功率: {successful_batch}/{len(batch_results)}")
        
        # 显示输出目录
        outputs_dir = Path("outputs")
        if outputs_dir.exists():
            file_manager = FileManager()
            dir_info = file_manager.get_directory_size(outputs_dir)
            logger.info(f"输出目录大小: {dir_info['total_size_mb']:.2f} MB")
            logger.info(f"生成文件数量: {dir_info['file_count']}")
        
    except Exception as e:
        logger.error(f"演示过程中出现错误: {e}")
        raise


if __name__ == "__main__":
    main()